#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <dirent.h>
#include <sys/stat.h>
#include <limits.h>
#include <time.h>

#define MAX_DICT_SIZE 4096
#define MAX_CODE_LEN 12
#define FILE_HEADER_MAGIC "RLEZIP"
#define FILE_HEADER_VERSION 1
#define MAX_PATH_LENGTH 4096

typedef unsigned short code_t;

typedef struct {
    char *sequence;
    int length;
} DictionaryEntry;

DictionaryEntry dictionary[MAX_DICT_SIZE];
int next_code = 256;

// 全局变量用于统计
long total_input_size = 0;
long total_output_size = 0;
time_t start_time, end_time;

void init_dictionary() {
	int i=0; 
    for ( i = 0; i < 256; i++) {
        dictionary[i].sequence = malloc(1);
        dictionary[i].sequence[0] = (char)i;
        dictionary[i].length = 1;
    }
    for ( i = 256; i < MAX_DICT_SIZE; i++) {
        dictionary[i].sequence = NULL;
        dictionary[i].length = 0;
    }
}

void free_dictionary() {
	int i=0;
    for ( i = 0; i < MAX_DICT_SIZE; i++) {
        if (dictionary[i].sequence != NULL) {
            free(dictionary[i].sequence);
        }
    }
}

void write_code(FILE *output, code_t code, int *bit_offset, unsigned char *buffer) {
	int i=0;
    for ( i = 0; i < MAX_CODE_LEN; i++) {
        if (code & (1 << i)) {
            *buffer |= (1 << *bit_offset);
        }
        (*bit_offset)++;
        if (*bit_offset == 8) {
            fwrite(buffer, 1, 1, output);
            total_output_size++;
            *buffer = 0;
            *bit_offset = 0;
        }
    }
}

void flush_buffer(FILE *output, int *bit_offset, unsigned char *buffer) {
    if (*bit_offset > 0) {
        fwrite(buffer, 1, 1, output);
        total_output_size++;
        *buffer = 0;
        *bit_offset = 0;
    }
}

void print_stats(const char *filename, long input_size, long output_size) {
    double ratio = (input_size == 0) ? 0.0 : (1.0 - (double)output_size / input_size) * 100.0;
    printf("%-40s %10ld -> %10ld bytes (%5.2f%% compression)\n", 
           filename, input_size, output_size, ratio);
}

void compress_file(FILE *input, FILE *output, const char *filename, int is_first_file) {
    static int dictionary_initialized = 0;
    if (!dictionary_initialized) {
        init_dictionary();
        dictionary_initialized = 1;
    }

    // 记录文件开始位置以计算压缩后大小
    long output_start_pos = ftell(output);

    // Write file header
    if (is_first_file) {
        fwrite(FILE_HEADER_MAGIC, 1, strlen(FILE_HEADER_MAGIC), output);
        unsigned char version = FILE_HEADER_VERSION;
        fwrite(&version, 1, 1, output);
        output_start_pos += strlen(FILE_HEADER_MAGIC) + 1;
    }

    // Write filename length and filename
    unsigned short filename_len = (unsigned short)strlen(filename);
    fwrite(&filename_len, sizeof(filename_len), 1, output);
    fwrite(filename, 1, filename_len, output);

    // Write file size
    fseek(input, 0, SEEK_END);
    long file_size = ftell(input);
    total_input_size += file_size;
    fseek(input, 0, SEEK_SET);
    fwrite(&file_size, sizeof(file_size), 1, output);

    // Compress data
    unsigned char current_char;
    char *current_sequence = NULL;
    int current_sequence_len = 0;
    int bit_offset = 0;
    unsigned char buffer = 0;

    while (fread(&current_char, 1, 1, input)) {
        char *new_sequence = realloc(current_sequence, current_sequence_len + 1);
        if (!new_sequence) {
            perror("Memory allocation failed");
            exit(EXIT_FAILURE);
        }
        current_sequence = new_sequence;
        current_sequence[current_sequence_len++] = current_char;

        int found = -1;
        int i=0;
        for ( i = 0; i < next_code; i++) {
            if (dictionary[i].length == current_sequence_len && 
                memcmp(dictionary[i].sequence, current_sequence, current_sequence_len) == 0) {
                found = i;
                break;
            }
        }

        if (found == -1) {
            if (next_code < MAX_DICT_SIZE) {
                dictionary[next_code].sequence = malloc(current_sequence_len);
                memcpy(dictionary[next_code].sequence, current_sequence, current_sequence_len);
                dictionary[next_code].length = current_sequence_len;
                next_code++;
            }

            // Write the code for the sequence without the last character
            char *sequence_to_write = current_sequence;
            int sequence_to_write_len = current_sequence_len - 1;
            int i=0;
            for ( i = 0; i < next_code; i++) {
                if (dictionary[i].length == sequence_to_write_len && 
                    memcmp(dictionary[i].sequence, sequence_to_write, sequence_to_write_len) == 0) {
                    write_code(output, (code_t)i, &bit_offset, &buffer);
                    break;
                }
            }

            // Start new sequence with the last character
            free(current_sequence);
            current_sequence = malloc(1);
            current_sequence[0] = current_char;
            current_sequence_len = 1;
        }
    }

    // Write the last sequence
    if (current_sequence_len > 0) {
    	int i=0;
        for ( i = 0; i < next_code; i++) {
            if (dictionary[i].length == current_sequence_len && 
                memcmp(dictionary[i].sequence, current_sequence, current_sequence_len) == 0) {
                write_code(output, (code_t)i, &bit_offset, &buffer);
                break;
            }
        }
    }

    flush_buffer(output, &bit_offset, &buffer);
    free(current_sequence);

    // 计算并显示当前文件的压缩统计
    long output_end_pos = ftell(output);
    long compressed_size = output_end_pos - output_start_pos;
    print_stats(filename, file_size, compressed_size);
}

void process_directory(const char *dir_path, const char *base_path, FILE *output, int *is_first_file) {
    DIR *dir;
    struct dirent *entry;
    struct stat statbuf;
    char full_path[MAX_PATH_LENGTH];
    char relative_path[MAX_PATH_LENGTH];

    if ((dir = opendir(dir_path)) == NULL) {
        perror("opendir failed");
        return;
    }

    while ((entry = readdir(dir)) != NULL) {
        if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) {
            continue;
        }

        snprintf(full_path, sizeof(full_path), "%s/%s", dir_path, entry->d_name);
        if (base_path[0] == '\0') {
            snprintf(relative_path, sizeof(relative_path), "%s", entry->d_name);
        } else {
            snprintf(relative_path, sizeof(relative_path), "%s/%s", base_path, entry->d_name);
        }

        if (stat(full_path, &statbuf) == -1) {
            perror("stat failed");
            continue;
        }

        if (S_ISDIR(statbuf.st_mode)) {
            process_directory(full_path, relative_path, output, is_first_file);
        } else {
            FILE *input = fopen(full_path, "rb");
            if (!input) {
                perror("fopen failed");
                continue;
            }
            compress_file(input, output, relative_path, *is_first_file);
            fclose(input);
            *is_first_file = 0;
        }
    }
    closedir(dir);
}

void print_summary() {
    end_time = time(NULL);
    double elapsed_time = difftime(end_time, start_time);
    double ratio = (total_input_size == 0) ? 0.0 : 
                  (1.0 - (double)total_output_size / total_input_size) * 100.0;
    
    printf("\nSummary:\n");
    printf("Total input size:  %15ld bytes\n", total_input_size);
    printf("Total output size: %15ld bytes\n", total_output_size);
    printf("Compression ratio: %15.2f%%\n", ratio);
    printf("Processing time:   %15.2f seconds\n", elapsed_time);
    printf("Compression speed: %15.2f KB/s\n", 
           (total_input_size / 1024.0) / (elapsed_time > 0 ? elapsed_time : 1));
}

int main(int argc, char *argv[]) {
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <input1> [input2 ...] <output>\n", argv[0]);
        return EXIT_FAILURE;
    }

    // 初始化统计
    total_input_size = 0;
    total_output_size = 0;
    start_time = time(NULL);

    FILE *output = fopen(argv[argc - 1], "wb");
    if (!output) {
        perror("fopen failed");
        return EXIT_FAILURE;
    }

    printf("%-40s %12s %12s %15s\n", "Filename", "Input Size", "Output Size", "Compression Ratio");
    printf("------------------------------------------------------------------------------\n");

    int is_first_file = 1;
    int i=0;
    for ( i = 1; i < argc - 1; i++) {
        struct stat statbuf;
        if (stat(argv[i], &statbuf) == -1) {
            perror("stat failed");
            continue;
        }

        if (S_ISDIR(statbuf.st_mode)) {
            process_directory(argv[i], "", output, &is_first_file);
        } else {
            FILE *input = fopen(argv[i], "rb");
            if (!input) {
                perror("fopen failed");
                continue;
            }
            compress_file(input, output, argv[i], is_first_file);
            fclose(input);
            is_first_file = 0;
        }
    }

    // Write end marker
    write_code(output, MAX_DICT_SIZE - 1, &(int){0}, &(unsigned char){0});
    flush_buffer(output, &(int){0}, &(unsigned char){0});

    fclose(output);
    free_dictionary();

    // 打印总结信息
    print_summary();

    return EXIT_SUCCESS;
}